#!/usr/bin/env python3
'''
未整改
'''
import ptan
import gymnasium as gym
import argparse
import numpy as np
import torch
import torch.nn as nn
import torch.nn.utils as nn_utils
import torch.nn.functional as F
from lib import common

import torch
import torch.nn.functional as F
import ale_py

gym.register_envs(ale_py)
def wrap_dqn(env, stack_frames=4, episodic_life=True, reward_clipping=True):
    # if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        # env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    # env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)
    # 跳帧包装器
    # env = ptan.common.wrappers.MaxAndSkipEnv(env, skip=4)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = ptan.common.wrappers.FireResetEnv(env)
    env = ptan.common.wrappers.ProcessFrame84(env)
    env = ptan.common.wrappers.ImageToPyTorch(env)
    env = ptan.common.wrappers.FrameStack(env, stack_frames)
    return env


class AtariA2C(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(AtariA2C, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(input_shape)
        self.policy = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, n_actions)
        )

        self.value = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU(),
            nn.Linear(512, 1)
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        fx = x.float() / 256
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        return self.policy(conv_out), self.value(conv_out)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("-m", "--model", required=True, help="Model file name")
    parser.add_argument("--cuda", default=False, action="store_true")
    parser.add_argument("--seed", type=int, default=0, help="Random seed")
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    torch.manual_seed(args.seed)
    np.random.seed(args.seed)

    make_env = lambda: wrap_dqn(gym.make("ALE/Assault-v5", render_mode="human"))
    env = make_env()
    net = AtariA2C(env.observation_space.shape, env.action_space.n)
    checkpoint = torch.load(args.model, map_location=device, weights_only=False)
    net.load_state_dict(checkpoint)
    if args.cuda:
        net.cuda()

    act_selector = ptan.actions.ArgmaxActionSelector()

    obs, _ = env.reset()
    total_reward = 0.0
    total_steps = 0

    while True:
        obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
        logits_v, values_v = net(obs_v)
        probs_v = F.softmax(logits_v, dim=1)
        probs = probs_v.data.cpu().numpy()
        actions = act_selector(probs)
        obs, r, done, trunc, info = env.step(actions[0])
        if r != 0:
            print("reward: %d" % r)
        total_reward += r
        total_steps += 1
        if done or trunc:
            print(info)
            break

    print("Done in %d steps, reward %.2f" % (total_steps, total_reward))
